import torch.nn as nn


class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        # 卷积层 1~3
        self.conv1 = nn.Sequential(
            nn.Conv2d(  # (1, 28, 28), 卷积后: (32, 28, 28), 池化后: (32, 14, 14)
                in_channels=1,  # 输入图片的层数, 灰色为 1 层, RGB 为 3 层
                out_channels=32,  # 输出图片的层数
                kernel_size=3,  # 卷积核大小: 3 * 3
                stride=1,  # 卷积核滑动步长: 1
                padding=1  # 图像周围填充 1 个元素, 保持卷积后特征图大小不变
            ),
            nn.ReLU(),  # 激活函数
            nn.MaxPool2d(kernel_size=2)  # 最大池化, 池化窗口大小: 2 * 2
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(  # (32, 14, 14), 卷积后: (64, 14, 14), 池化后: (64, 7, 7)
                in_channels=32,  # 此处输入图片的层数是上层输出图片的层数
                out_channels=64,
                kernel_size=3,
                stride=1,
                padding=1
            ),
            nn.ReLU(),  # 激活函数
            nn.MaxPool2d(kernel_size=2)  # 最大池化, 池化窗口大小: 2 * 2
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(  # (64, 7, 7), 卷积后: (128, 7, 7), 池化后: (128, 3, 3)
                in_channels=64,
                out_channels=128,
                kernel_size=3,
                stride=1,
                padding=1
            ),
            nn.ReLU(),  # 激活函数
            nn.MaxPool2d(kernel_size=2)
        )
        # 全连接层 1~3
        self.fc1 = nn.Linear(128 * 3 * 3, 256)
        self.fc2 = nn.Linear(256, 128)
        self.fc3 = nn.Linear(128, 10)

    def forward(self, x):
        x = self.conv1(x)  # 通过第一个卷积层
        x = self.conv2(x)  # 通过第二个卷积层
        x = self.conv3(x)  # 通过第三个卷积层
        x = x.view(x.size(0), -1)  # 展平操作, 将多维数据转换为二维数据
        x = self.fc1(x)  # 通过第一个全连接层
        x = self.fc2(x)  # 通过第二个全连接层
        output = self.fc3(x)  # 通过第三个全连接层
        return output
